package com.lagou.niancheng.lagouservicegateway;

import com.google.common.util.concurrent.RateLimiter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.route.RouteLocator;
import org.springframework.cloud.gateway.route.builder.RouteLocatorBuilder;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * com.lagou.niancheng.lagouservicegateway.AccessTokenFilter
 *
 * @author 大犬
 * @date 2020/7/1
 */
@Configuration
public class IpLimitFilterConf {

    private static final Logger logger = LoggerFactory.getLogger(IpLimitFilterConf.class);

    @Value("${lagou.gateway.server.minutes}")
    private long minutes;
    @Value("${lagou.gateway.server.limit}")
    private long limit;


    @Bean
    public RouteLocator ipRouteLocator(RouteLocatorBuilder builder) {

        return builder.routes()
                .route(r -> r.path("/api/user/register")
                        .filters(f -> f.stripPrefix(1).filter(new IpLimitFilter((double) limit / (double) minutes / 60D)))
                        .uri("lb://user-server")
                        .order(0)
                        .id("token_filter_router"))
                .build();
    }

    public static class IpLimitFilter implements GatewayFilter {

        private double perSecondsLimit;

        public IpLimitFilter(double perSecondsLimit) {
            this.perSecondsLimit = perSecondsLimit;
        }

        private Map<String, RateLimiter> rateLimiterMap = new ConcurrentHashMap<>();

        @Override
        public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
            ServerHttpRequest request = exchange.getRequest();
            ServerHttpResponse response = exchange.getResponse();
            String hostString = request.getRemoteAddress().getHostString();
            rateLimiterMap.putIfAbsent(hostString, RateLimiter.create(perSecondsLimit));

            RateLimiter ipLimiter = rateLimiterMap.get(hostString);
            if (ipLimiter.tryAcquire()) {
                return chain.filter(exchange);
            } else {
                return noExceedLimit(response);
            }
        }

        private Mono<Void> noExceedLimit(ServerHttpResponse response) {
            //设置status
            logger.info("超出限制");
            response.setStatusCode(HttpStatus.SEE_OTHER);
            response.getHeaders().add("Content-Type", "text/html;charset=UTF-8");
            byte[] bytes = "您频繁进⾏注册，请求已被拒绝".getBytes(StandardCharsets.UTF_8);
            DataBuffer buffer = response.bufferFactory().wrap(bytes);
            return response.writeWith(Flux.just(buffer));
        }
    }
}
